# encoding:utf-8
'''
Created on 2020-12-23
@author: adog
ref: https://scikit-learn.org/stable/auto_examples/applications/plot_stock_market.html#sphx-glr-auto-examples-applications-plot-stock-market-py
'''
import apps.adog.dataMan as DM
import util.xtool as xtool
import util.tradeDate as TD
import apps.adog.util.dbcom as dbcom
import apps.adog.util.mdbcom as mdbcom

import numpy as np
import pandas as pd
import apps.adog.codStartDate as csd
codStartDate=csd.codStartDate



def getColor(v):
    color=[
            '#285818', #<-0.05
            '#4ca22f', #-0.05 -0.02
            '#8de071', #-0.02 -0.005
            '#d0fbcb', # -0.005 0.
            '#edd9dc', # 0 0,005
            '#e27485', # 0,005 0.02
            '#e1604c', # 0.02 0.05
            '#822a1c', #>0.05
           ]
           
    c='#FFFFFF'
    threshold=[-0.05,-0.02,-0.005,0., 0.005,0.02,0.05]
    if v<=threshold[0]:
        c=color[0]
    elif v>threshold[-1]:
        c=color[-1]
    else:
        for i in range(len(threshold)-1):
            if (v>threshold[i])&(v<=threshold[i+1]):
                c=color[i+1]
    return c
    
class Relation():
    def __init__(self,startDate=None, endDate=None,tradeDate=None,freq='1d', is15minToday=False,vs_info=None,obs_exclusive=[]):
        self.field=None
        obs_exclusive=obs_exclusive
        self.is15minToday=is15minToday
        self.data=None
        self.futureInfo=None
        self.tradeDate=tradeDate
        self.todayTime=None
        self.vs_info=vs_info
        
        self.loadData(startDate, endDate, freq, is15minToday, obs_exclusive)
        self.setBlkIdx()
        
    def loadData(self,startDate=None, endDate=None,freq=None, is15minToday=False,obs_exclusive=None): 
        if is15minToday:
            self.field=['date','time', 'cod','code_cn','chg_cls_r','chg_vol_r']
            if 'vs prev day' in self.vs_info:
                dm_prevDay=DM.DataMan(startDate=TD.tradeDateBefAftX(xtool.nowDate(), -20),endDate=xtool.nowDate(),freq='1d',is15minToday=False)
                data_prevDay=dm_prevDay.data[dm_prevDay.data['date']==dm_prevDay.data['date'].max()]
                
                dm=DM.DataMan(is15minToday=True)
                data_realtime=pd.merge(left=dm.data[['date','time', 'cod','code_cn','close','volume']],\
                                       right=data_prevDay[['cod','close','volume']],\
                                       on='cod',\
                                       how='left')
                
                data_realtime['chg_cls_r']=data_realtime['close_x']/data_realtime['close_y']-1.
                data_realtime['chg_vol_r']=data_realtime['volume_x']/data_realtime['volume_y']-1.
                
                self.data=data_realtime.replace([np.inf, -np.inf], 0.)[self.field]
                self.futureInfo=dm.future_info.copy()[['_id','BLK','blk_cn']]
                self.todayTime=sorted(self.data['time'].unique()[self.data['time'].unique()<xtool.nowOnlyTimeMins()].tolist())
            elif 'vs prev 15min' in self.vs_info:
                dm=DM.DataMan(is15minToday=True)
                self.data=dm.data
                self.futureInfo=dm.future_info.copy()[['_id','BLK','blk_cn']]
                self.todayTime=sorted(self.data['time'].unique()[self.data['time'].unique()<xtool.nowOnlyTimeMins()].tolist())
                
        else:
            if self.tradeDate is None:    
                sDate=startDate
                eDate=xtool.nowDate()
                if endDate is not None:
                    eDate=endDate
                    
                self.startDate=sDate
                self.endDate=eDate
                tradeDate=TD.tradeDateDf
                self.tradeDate=tradeDate.loc[(tradeDate['tradeDate']>=sDate)&(tradeDate['tradeDate']<=eDate)]
                self.tradeDate.reset_index(drop=True, inplace=True)    
                self.tradeDate=self.tradeDate['tradeDate'].tolist()
                
            else:
                sDate= self.tradeDate[0]
                eDate= self.tradeDate[-1]
                
            self.field=['date','cod','code_cn','chg_cls_r','chg_vol_r']
           
            obs=sorted([k for k,v in codStartDate.items() if (v<=sDate)&(not k in obs_exclusive)])
            dm=DM.DataMan(startDate=sDate,endDate=eDate,freq=freq,obs=obs,is15minToday=False)
            self.data=dm.data
            self.futureInfo=dm.future_info.copy()[['_id','BLK','blk_cn']]
          
        
        self.data.sort_values(by='cod',ascending=1,inplace=True)
        if is15minToday:
            self.data.drop_duplicates(['date','time','cod'], keep='last', inplace=True)
        else:
            self.data.drop_duplicates(['date','cod'], keep='last', inplace=True)
        self.data.reset_index(drop=True,inplace=True)
        #data=pd.read_csv('./data.csv')
        #self.futureInfo=pd.read_csv('./fi.csv')[['_id','BLK','blk_cn']] 
        
    def setBlkIdx(self):     
        self.futureInfo['bb']=[self.futureInfo.loc[i,'blk_cn'].split('_')[0] for i in self.futureInfo.index]
        self.futureInfo['kk']=[self.futureInfo.loc[i,'blk_cn'].split('_')[1] for i in self.futureInfo.index] 
        idx_bb,idx_kk={},{}
        fbb=self.futureInfo['bb'].unique()
        for i in range(len(fbb)):
            idx_bb[fbb[i]]=i 
            
        fkk=self.futureInfo['kk'].unique()
        for j in range(len(fkk)):
            idx_kk[fkk[j]]=j 
        self.futureInfo['bb_idx']=[idx_bb[self.futureInfo.loc[i, 'bb']] for i in self.futureInfo.index]
        self.futureInfo['kk_idx']=[idx_kk[self.futureInfo.loc[i, 'kk']] for i in self.futureInfo.index]
        
    def getCorr(self,startDate,endDate):
        data=self.data.copy().pivot(index='date',columns='cod')
        corr=data['chg_cls_r'].copy().astype(float).fillna(0.0).corr()
        corr=corr.stack()
        corr.index.names=['cod1','cod2']
        corr=corr.reset_index(drop=False)
        corr.rename(columns={0:'corr'},inplace=True)
        corr['startDate']=startDate
        corr['endDate']=endDate
        corr['period']=TD.getTradeDayDiff(startDate, endDate)
        corr['_id']=startDate+'_'+endDate+'_'+corr['cod1']+'_'+corr['cod2']
        return corr
        #mdbcom.saveBatch('corr', corr.to_dict(orient='records'))
        #dbcom.creatTableReplaceMany('corr', corr, drop_exist=False)
    
    def getNodes(self):
        #data=self.data.copy()
        data=pd.merge(left=self.data,right=self.futureInfo, left_on='cod',right_on='_id',how='left')
        nodeTimeList=[]
        dtime=None
        
        if self.is15minToday:
            dtime=self.todayTime
        else:
            dtime=self.tradeDate
        
        for dt in dtime:
            di=None
            if self.is15minToday:
                di=data.loc[data['time']==dt,self.field+['bb','bb_idx']]
            else:
                di=data.loc[data['date']==dt,self.field+['bb','bb_idx']]
                
            nodeList=[]
            for i in di.index:
                node={
                      'id': di.loc[i,'cod'],
                      'label':di.loc[i,'code_cn'].replace('期货指数','')+di.loc[i,'cod'],
                      'size':int(abs(di.loc[i,'chg_cls_r'])*10000),#(100 if self.is15minToday else 10000)), 
                      'chg_cls_r':round(di.loc[i,'chg_cls_r'],4),
                      'comboId':di.loc[i,'bb'],
                      'cluster':di.loc[i,'bb_idx'],
                      'style':{
                               'fill': getColor(di.loc[i,'chg_cls_r']),
                               }
                      }
                nodeList.append(node)
                
            '''    
            codLink=pd.read_csv('codLink.txt')
            newCod=sorted(set([c for c in codLink['cod1'].tolist()+codLink['cod2'].tolist() if c[0]=='+']))
            for n in newCod:
                newNode={
                      'id': n,
                      'label':n,
                      'size':50, 
                    
                      }
                nodeList.append(newNode)
            '''
                
            if len(nodeList)>0: 
                nodeTimeList.append([dt,nodeList])
        
        return nodeTimeList
                            
    def getEdges(self):
        
        codLink=pd.read_csv(xtool.adogRootPath()+'/data/codLink.txt')
        edgeList=[]
        for i in codLink.index:
            edge={
                   'id':codLink.loc[i,'cod1']+'-'+codLink.loc[i,'cod2'],
                  #'label':codLink.loc[i,'cod1']+'-'+codLink.loc[i,'cod2'], 
                  'source':codLink.loc[i,'cod1'],
                  'target':codLink.loc[i,'cod2'],                     
                  }
            edgeList.append(edge)
            
        return edgeList
    
    def getCombos(self):
        #fi=self.futureInfo.copy()[['_id','BLK','blk_cn']]
        comboList=[]
        for b in self.futureInfo['bb'].unique():
            combo={
                   'id': b,
                   'label': b,
                   
                   }
            comboList.append(combo)
        '''    
        for i in self.futureInfo.index:
            
            combop={
                   'id': self.futureInfo.loc[i,'kk'],
                   'label': self.futureInfo.loc[i,'kk'],
                   'parentId':self.futureInfo.loc[i,'bb'],
                   
                   }
            comboList.append(combop)
        '''
        return comboList
    
    def getHulls(self):
        fi=self.futureInfo.copy()[['_id','BLK','blk_cn']]
        hullList=[]
        for b in fi['BLK'].unique():
            hull={
                   'id': b,
                   'label': b,
                   'member':fi.loc[fi['BLK']==b,'_id'].tolist()
                   }
            hullList.append(hull)
        return hullList
    
if __name__=='__main__':
    r=Relation()
    #r.getCorr(startDate='2020-10-01', endDate='2021-03-01')
    n=r.getNodes('2021-04-02')
    e=r.getEdges()
    
